from typing import Dict, Any, Optional
from src.core.output_manager_base import OutputManagerBase, SimulationParameters, MPI, fem, FENICSX_AVAILABLE, Mesh, plot

class PyVistaVisualizer(OutputManagerBase):
    """
    使用 PyVista 进行实时或最终可视化（如果适用）。
    注意：实时可视化在复杂的 HPC 环境中可能不实用或不被推荐。
    此类的行为可能高度依赖于运行环境。
    """
    def __init__(self, 
                 comm: Optional[MPI.Comm], 
                 params: SimulationParameters, 
                 fem_entities: Dict[str, Any], 
                 output_dir: str,
                 visualize_frequency: int = 10,
                 interactive_visualization: bool = False):
        """
        初始化 PyVistaVisualizer。

        参数:
            comm (Optional[MPI.Comm]): MPI 通信器。
            params (SimulationParameters): 模拟参数。
            fem_entities (Dict[str, Any]): FEM 实体。
            output_dir (str): 输出目录 (可能用于保存截图)。
            visualize_frequency (int): 可视化频率。
            interactive_visualization (bool): 是否尝试交互式可视化。
        """
        super().__init__(comm, params, fem_entities, output_dir)
        self.visualize_frequency = visualize_frequency
        self.interactive_visualization = interactive_visualization
        self.plotter = None # PyVista Plotter 对象
        
        if self.comm is None or self.comm.rank == 0: # type: ignore
            print(f"PyVistaVisualizer: 初始化。可视化频率: {self.visualize_frequency}, 交互式: {self.interactive_visualization}。")
            if self.interactive_visualization:
                 print("警告: PyVista 的交互式可视化可能不适用于所有环境。")

    def is_output_step(self, iteration_count: int, current_time: float) -> bool:
        """
        判断是否为可视化步骤。
        """
        # 仅在主进程上执行可视化
        if self.comm is not None and self.comm.rank != 0: # type: ignore
            return False
        return iteration_count % self.visualize_frequency == 0

    def execute_step(self, iteration_count: int, current_time: float) -> None:
        """
        执行可视化步骤。
        """
        if self.comm is not None and self.comm.rank != 0: # type: ignore
            return
            
        print(f"PyVistaVisualizer: 在迭代 {iteration_count}, 时间 {current_time:.2f} 尝试执行可视化。")
        if not FENICSX_AVAILABLE:
            print("  PyVistaVisualizer: FEniCSx 不可用，跳过可视化。")
            return

        try:
            import pyvista # type: ignore
            # plot 已经从 output_manager_base 导入

            # 示例：可视化第一个找到的 fem.Function
            function_to_plot = None
            mesh_to_plot = self.fem_entities.get("mesh", None)
            # 使用从 output_manager_base 导入的 Mesh 类型进行检查
            if not isinstance(mesh_to_plot, Mesh): # type: ignore
                print(f"  PyVistaVisualizer: fem_entities 中未找到有效的 Mesh 对象。类型为: {type(mesh_to_plot)}")
                return


            for key, entity in self.fem_entities.items():
                if isinstance(entity, fem.Function): # type: ignore
                    function_to_plot = entity
                    print(f"  PyVistaVisualizer: 找到函数 \'{entity.name if hasattr(entity, 'name') else key}\' 进行可视化。") # type: ignore
                    break
            
            if function_to_plot is None:
                print("  PyVistaVisualizer: 未找到可供可视化的 fem.Function。")
                return
            
            # 获取拓扑和几何信息
            # VTXWriter 的 plot.vtk_mesh 可能是一个好的参考
            # 如果要使用 plot.create_vtk_mesh，它已经从 output_manager_base 导入
            # cells, types, x = plot.create_vtk_mesh(function_to_plot.function_space) # 使用导入的 plot
            # grid = pyvista.UnstructuredGrid(cells, types, x)
            # grid.point_data["u"] = function_to_plot.x.array.real 
            # grid.set_active_scalars("u")
            
            # 简化：直接使用 PyVista 的 dolfinx 兼容功能 (如果存在，或自己实现转换)
            # pyvista.Plotter() 等。为保持骨架简单，此处仅打印消息。
            print(f"  PyVistaVisualizer: (模拟) 正在使用 PyVista 可视化函数 \'{function_to_plot.name if hasattr(function_to_plot, 'name') else 'unknown_function'}\'。") # type: ignore
            
            # if self.interactive_visualization:
            #     if self.plotter is None:
            #         self.plotter = pyvista.Plotter(notebook=False) # notebook=False for standalone window
            #     self.plotter.clear()
            #     self.plotter.add_mesh(grid, show_edges=True)
            #     self.plotter.show(auto_close=False) # or use plotter.render() in a loop for non-blocking
            # else:
            #     # 保存截图
            #     # plotter_screenshot = pyvista.Plotter(off_screen=True)
            #     # plotter_screenshot.add_mesh(grid)
            #     # screenshot_path = os.path.join(self.output_dir, f"visualization_{iteration_count}.png")
            #     # plotter_screenshot.screenshot(screenshot_path)
            #     # print(f"  PyVistaVisualizer: 截图已保存到 {screenshot_path}")
            pass

        except ImportError: # 主要捕获 pyvista 的导入错误
            print("警告: PyVistaVisualizer 无法导入 pyvista，无法进行可视化。")
        except Exception as e:
            print(f"错误: PyVistaVisualizer 可视化时发生错误: {e}")

    def finalize(self) -> None:
        """
        完成可视化操作。
        """
        if self.comm is not None and self.comm.rank != 0: # type: ignore
            return

        if self.plotter is not None:
            self.plotter.close()
            print("PyVistaVisualizer: PyVista Plotter 已关闭。")
        print("PyVistaVisualizer: 完成所有可视化操作。") 